# type: ignore
"""first migration

Revision ID: c6f47d7d2fa1
Revises:
Create Date: 2025-01-10 19:06:15.190427+00:00

"""

from typing import Sequence, Union

import sqlalchemy as sa
from alembic import op
from pgvector.sqlalchemy import Vector
from sqlalchemy.dialects import postgresql

# revision identifiers, used by Alembic.
revision: str = "c6f47d7d2fa1"
down_revision: Union[str, None] = None
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # create extension if not exists vector;
    op.execute("CREATE EXTENSION IF NOT EXISTS vector;")

    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table(
        "apps",
        sa.Column("id", sa.UUID(), nullable=False),
        sa.Column("name", sa.String(length=100), nullable=False),
        sa.Column("display_name", sa.String(length=255), nullable=False),
        sa.Column("provider", sa.String(length=255), nullable=False),
        sa.Column("version", sa.String(length=255), nullable=False),
        sa.Column("description", sa.Text(), nullable=False),
        sa.Column("logo", sa.Text(), nullable=True),
        sa.Column("categories", postgresql.ARRAY(sa.String()), nullable=False),
        sa.Column("visibility", sa.Enum("PUBLIC", "PRIVATE", name="visibility"), nullable=False),
        sa.Column("active", sa.Boolean(), nullable=False),
        sa.Column("security_schemes", sa.JSON(), nullable=False),
        sa.Column("default_security_credentials_by_scheme", sa.JSON(), nullable=False),
        sa.Column("embedding", Vector(dim=1024), nullable=False),
        sa.Column("created_at", sa.DateTime(), server_default=sa.text("now()"), nullable=False),
        sa.Column("updated_at", sa.DateTime(), server_default=sa.text("now()"), nullable=False),
        sa.PrimaryKeyConstraint("id"),
        sa.UniqueConstraint("name"),
    )
    op.create_table(
        "entities",
        sa.Column("id", sa.UUID(), nullable=False),
        sa.Column(
            "type", sa.Enum("ENTITY", "USER", "ORGANIZATION", name="entitytype"), nullable=False
        ),
        sa.Column("name", sa.String(length=255), nullable=False),
        sa.Column("email", sa.String(length=255), nullable=False),
        sa.Column("profile_picture", sa.Text(), nullable=True),
        sa.Column("created_at", sa.DateTime(), server_default=sa.text("now()"), nullable=False),
        sa.Column("updated_at", sa.DateTime(), server_default=sa.text("now()"), nullable=False),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_table(
        "functions",
        sa.Column("id", sa.UUID(), nullable=False),
        sa.Column("app_id", sa.UUID(), nullable=False),
        sa.Column("name", sa.String(length=255), nullable=False),
        sa.Column("description", sa.Text(), nullable=False),
        sa.Column("tags", postgresql.ARRAY(sa.String()), nullable=False),
        sa.Column("visibility", sa.Enum("PUBLIC", "PRIVATE", name="visibility"), nullable=False),
        sa.Column("active", sa.Boolean(), nullable=False),
        sa.Column("protocol", sa.Enum("REST", name="protocol"), nullable=False),
        sa.Column("protocol_data", sa.JSON(), nullable=False),
        sa.Column("parameters", sa.JSON(), nullable=False),
        sa.Column("response", sa.JSON(), nullable=False),
        sa.Column("embedding", Vector(dim=1024), nullable=False),
        sa.Column("created_at", sa.DateTime(), server_default=sa.text("now()"), nullable=False),
        sa.Column("updated_at", sa.DateTime(), server_default=sa.text("now()"), nullable=False),
        sa.ForeignKeyConstraint(
            ["app_id"],
            ["apps.id"],
        ),
        sa.PrimaryKeyConstraint("id"),
        sa.UniqueConstraint("name"),
    )
    op.create_table(
        "organizations",
        sa.Column("id", sa.UUID(), nullable=False),
        sa.ForeignKeyConstraint(
            ["id"],
            ["entities.id"],
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_table(
        "projects",
        sa.Column("id", sa.UUID(), nullable=False),
        sa.Column("owner_id", sa.UUID(), nullable=False),
        sa.Column("name", sa.String(length=255), nullable=False),
        sa.Column(
            "visibility_access", sa.Enum("PUBLIC", "PRIVATE", name="visibility"), nullable=False
        ),
        sa.Column("daily_quota_used", sa.Integer(), nullable=False),
        sa.Column(
            "daily_quota_reset_at", sa.DateTime(), server_default=sa.text("now()"), nullable=False
        ),
        sa.Column("total_quota_used", sa.Integer(), nullable=False),
        sa.Column("created_at", sa.DateTime(), server_default=sa.text("now()"), nullable=False),
        sa.Column("updated_at", sa.DateTime(), server_default=sa.text("now()"), nullable=False),
        sa.ForeignKeyConstraint(
            ["owner_id"],
            ["entities.id"],
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_table(
        "subscriptions",
        sa.Column("id", sa.UUID(), nullable=False),
        sa.Column("entity_id", sa.UUID(), nullable=False),
        sa.Column(
            "plan",
            sa.Enum("CUSTOM", "FREE", "PRO", "ENTERPRISE", name="subscriptionplan"),
            nullable=False,
        ),
        sa.Column(
            "status",
            sa.Enum("ACTIVE", "CANCELLED", "EXPIRED", name="subscriptionstatus"),
            nullable=False,
        ),
        sa.Column("expires_at", sa.DateTime(), nullable=True),
        sa.Column("created_at", sa.DateTime(), server_default=sa.text("now()"), nullable=False),
        sa.Column("updated_at", sa.DateTime(), server_default=sa.text("now()"), nullable=False),
        sa.ForeignKeyConstraint(
            ["entity_id"],
            ["entities.id"],
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_table(
        "users",
        sa.Column("id", sa.UUID(), nullable=False),
        sa.Column("identity_provider", sa.String(length=255), nullable=False),
        sa.Column("user_id_by_provider", sa.String(length=255), nullable=False),
        sa.ForeignKeyConstraint(
            ["id"],
            ["entities.id"],
        ),
        sa.PrimaryKeyConstraint("id"),
        sa.UniqueConstraint(
            "identity_provider", "user_id_by_provider", name="uc_auth_provider_user"
        ),
    )
    op.create_table(
        "agents",
        sa.Column("id", sa.UUID(), nullable=False),
        sa.Column("project_id", sa.UUID(), nullable=False),
        sa.Column("name", sa.String(length=255), nullable=False),
        sa.Column("description", sa.Text(), nullable=False),
        sa.Column("excluded_apps", postgresql.ARRAY(sa.UUID()), nullable=False),
        sa.Column("excluded_functions", postgresql.ARRAY(sa.UUID()), nullable=False),
        sa.Column("created_at", sa.DateTime(), server_default=sa.text("now()"), nullable=False),
        sa.Column("updated_at", sa.DateTime(), server_default=sa.text("now()"), nullable=False),
        sa.ForeignKeyConstraint(
            ["project_id"],
            ["projects.id"],
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_table(
        "app_configurations",
        sa.Column("id", sa.UUID(), nullable=False),
        sa.Column("project_id", sa.UUID(), nullable=False),
        sa.Column("app_id", sa.UUID(), nullable=False),
        sa.Column(
            "security_scheme",
            sa.Enum(
                "NO_AUTH", "API_KEY", "HTTP_BASIC", "HTTP_BEARER", "OAUTH2", name="securityscheme"
            ),
            nullable=False,
        ),
        sa.Column("security_scheme_overrides", sa.JSON(), nullable=False),
        sa.Column("enabled", sa.Boolean(), nullable=False),
        sa.Column("all_functions_enabled", sa.Boolean(), nullable=False),
        sa.Column("enabled_functions", postgresql.ARRAY(sa.UUID()), nullable=False),
        sa.Column("created_at", sa.DateTime(), server_default=sa.text("now()"), nullable=False),
        sa.Column("updated_at", sa.DateTime(), server_default=sa.text("now()"), nullable=False),
        sa.ForeignKeyConstraint(
            ["app_id"],
            ["apps.id"],
        ),
        sa.ForeignKeyConstraint(
            ["project_id"],
            ["projects.id"],
        ),
        sa.PrimaryKeyConstraint("id"),
        sa.UniqueConstraint("project_id", "app_id", name="uc_project_app"),
    )
    op.create_table(
        "linked_accounts",
        sa.Column("id", sa.UUID(), nullable=False),
        sa.Column("project_id", sa.UUID(), nullable=False),
        sa.Column("app_id", sa.UUID(), nullable=False),
        sa.Column("linked_account_owner_id", sa.String(length=255), nullable=False),
        sa.Column(
            "security_scheme",
            sa.Enum(
                "NO_AUTH", "API_KEY", "HTTP_BASIC", "HTTP_BEARER", "OAUTH2", name="securityscheme"
            ),
            nullable=False,
        ),
        sa.Column("security_credentials", sa.JSON(), nullable=False),
        sa.Column("enabled", sa.Boolean(), nullable=False),
        sa.Column("created_at", sa.DateTime(), server_default=sa.text("now()"), nullable=False),
        sa.Column("updated_at", sa.DateTime(), server_default=sa.text("now()"), nullable=False),
        sa.ForeignKeyConstraint(
            ["app_id"],
            ["apps.id"],
        ),
        sa.ForeignKeyConstraint(
            ["project_id"],
            ["projects.id"],
        ),
        sa.PrimaryKeyConstraint("id"),
        sa.UniqueConstraint(
            "project_id",
            "app_id",
            "linked_account_owner_id",
            name="uc_project_app_linked_account_owner",
        ),
    )
    op.create_table(
        "memberships",
        sa.Column("id", sa.UUID(), nullable=False),
        sa.Column("user_id", sa.UUID(), nullable=False),
        sa.Column("organization_id", sa.UUID(), nullable=False),
        sa.Column("role", sa.Enum("ADMIN", "MEMBER", name="organizationrole"), nullable=False),
        sa.Column("created_at", sa.DateTime(), server_default=sa.text("now()"), nullable=False),
        sa.Column("updated_at", sa.DateTime(), server_default=sa.text("now()"), nullable=False),
        sa.ForeignKeyConstraint(
            ["organization_id"],
            ["organizations.id"],
        ),
        sa.ForeignKeyConstraint(
            ["user_id"],
            ["users.id"],
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_table(
        "api_keys",
        sa.Column("id", sa.UUID(), nullable=False),
        sa.Column("key", sa.String(length=255), nullable=False),
        sa.Column("agent_id", sa.UUID(), nullable=False),
        sa.Column(
            "status", sa.Enum("ACTIVE", "DISABLED", "DELETED", name="apikeystatus"), nullable=False
        ),
        sa.Column("created_at", sa.DateTime(), server_default=sa.text("now()"), nullable=False),
        sa.Column("updated_at", sa.DateTime(), server_default=sa.text("now()"), nullable=False),
        sa.ForeignKeyConstraint(
            ["agent_id"],
            ["agents.id"],
        ),
        sa.PrimaryKeyConstraint("id"),
        sa.UniqueConstraint("agent_id"),
        sa.UniqueConstraint("key"),
    )
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_table("api_keys")
    op.drop_table("memberships")
    op.drop_table("linked_accounts")
    op.drop_table("app_configurations")
    op.drop_table("agents")
    op.drop_table("users")
    op.drop_table("subscriptions")
    op.drop_table("projects")
    op.drop_table("organizations")
    op.drop_table("functions")
    op.drop_table("entities")
    op.drop_table("apps")
    # ### end Alembic commands ###

    # Drop the Enum types
    op_bind = op.get_bind()
    sa.Enum(
        "PUBLIC",
        "PRIVATE",
        name="visibility",
    ).drop(op_bind)
    sa.Enum(
        "ENTITY",
        "USER",
        "ORGANIZATION",
        name="entitytype",
    ).drop(op_bind)
    sa.Enum(
        "ACTIVE",
        "DISABLED",
        "DELETED",
        name="apikeystatus",
    ).drop(op_bind)
    sa.Enum(
        "REST",
        name="protocol",
    ).drop(op_bind)
    sa.Enum(
        "CUSTOM",
        "FREE",
        "PRO",
        "ENTERPRISE",
        name="subscriptionplan",
    ).drop(op_bind)
    sa.Enum(
        "ACTIVE",
        "CANCELLED",
        "EXPIRED",
        name="subscriptionstatus",
    ).drop(op_bind)
    sa.Enum(
        "NO_AUTH",
        "API_KEY",
        "HTTP_BASIC",
        "HTTP_BEARER",
        "OAUTH2",
        name="securityscheme",
    ).drop(op_bind)
    sa.Enum(
        "ADMIN",
        "MEMBER",
        name="organizationrole",
    ).drop(op_bind)
    # drop extentions
    op.execute("DROP EXTENSION IF EXISTS vector;")
